{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4: Federated Learning with Model Averaging\n",
    "\n",
    "**Recap:** For Part 2 of this tutorial, we use simple version of Federated Learning to train a model. This step require sey to see gradient each data owner must to trust model owner.\n",
    "\n",
    "**Description:** For this tutorial, we go show you hiw you fit use the advanced aggregation tools wey dey Part 3 to allow weights to dey aggregated by trusted \"secure worker\" before we go send final resulting model to the model owner (wey be us).\n",
    "\n",
    "\n",
    "For this way wey we dey talk so, only secure go see who get the weights and where he come from. We go fit tell the part of the model wey we change, the problem na we no know which worker (Bob or Alice) change am, wey come create layer of privacy.\n",
    "\n",
    "Person wey write am:\n",
    "- Andrew Trask - Twitter: [@iamtrask](https://twitter.com/iamtrask)\n",
    "- Jason Mancuso - Twitter: [@jvmancuso](https://twitter.com/jvmancuso)\n",
    "\n",
    "Person wey translate am:\n",
    "- Temitọpẹ Ọladokun - Twitter: [@techie991](https://twitter.com/techie991)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import syft as sy\n",
    "import copy\n",
    "hook = sy.TorchHook(torch)\n",
    "from torch import nn, optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Make We Create Data Owners\n",
    "\n",
    "First tin wey we go do na to create two data owners (Bob and Alice) with small amount of data. We go initialize a secure machine called \"secure_worker\". In practice we fit say na secure hardware he be (sometin like Intel's SGX) or make we say na trusted intermediary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make we create couple workers\n",
    "\n",
    "bob = sy.VirtualWorker(hook, id=\"bob\")\n",
    "alice = sy.VirtualWorker(hook, id=\"alice\")\n",
    "secure_worker = sy.VirtualWorker(hook, id=\"secure_worker\")\n",
    "\n",
    "\n",
    "# A Toy Dataset\n",
    "data = torch.tensor([[0,0],[0,1],[1,0],[1,1.]], requires_grad=True)\n",
    "target = torch.tensor([[0],[0],[1],[1.]], requires_grad=True)\n",
    "\n",
    "# we go get pointers to training data on each worker if\n",
    "# we send some training data to bob and alice\n",
    "bobs_data = data[0:2].send(bob)\n",
    "bobs_target = target[0:2].send(bob)\n",
    "\n",
    "alices_data = data[2:].send(alice)\n",
    "alices_target = target[2:].send(alice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Make We Create Our Model\n",
    "\n",
    "For dis example, we go train with a simple Linear model. We go initialize am normally using Pytorch's nn.Linear constructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We go initialize A Toy Model\n",
    "model = nn.Linear(2,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Make we Send a Copy of the Model to Alice and Bob\n",
    "\n",
    "Next tin, na to send a copy of the current model to Alice and Bob so that they go fit perform steps of learning wey dem fit do on their own datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bobs_model = model.copy().send(bob)\n",
    "alices_model = model.copy().send(alice)\n",
    "\n",
    "bobs_opt = optim.SGD(params=bobs_model.parameters(),lr=0.1)\n",
    "alices_opt = optim.SGD(params=alices_model.parameters(),lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Make WeTrain Bob's and Alice's Models (in parallel)\n",
    "\n",
    "Normal tin for Federated Learning vaia Secure Averaging, data owner go first train their model for several iterations locally before we go average the models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "\n",
    "    # We go train Bob's Model\n",
    "    bobs_opt.zero_grad()\n",
    "    bobs_pred = bobs_model(bobs_data)\n",
    "    bobs_loss = ((bobs_pred - bobs_target)**2).sum()\n",
    "    bobs_loss.backward()\n",
    "\n",
    "    bobs_opt.step()\n",
    "    bobs_loss = bobs_loss.get().data\n",
    "\n",
    "    # We go train Alice's Model\n",
    "    alices_opt.zero_grad()\n",
    "    alices_pred = alices_model(alices_data)\n",
    "    alices_loss = ((alices_pred - alices_target)**2).sum()\n",
    "    alices_loss.backward()\n",
    "\n",
    "    alices_opt.step()\n",
    "    alices_loss = alices_loss.get().data\n",
    "    \n",
    "    print(\"Bob:\" + str(bobs_loss) + \" Alice:\" + str(alices_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 5: Make we Send Both Updated Models to a Secure Worker\n",
    "\n",
    "Now wey every data owner get partially trained model, we go average them in a secure way. We go achieve this by instructing Alice and Bob make dem send their model to the secure (trusted) server.\n",
    "\n",
    "Look am wella, as we go come use our API, he maen sey we go send each model DIRECTLY to the secure_worker. We never see am."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alices_model.move(secure_worker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bobs_model.move(secure_worker)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 6: Make We Average the Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the last step for this training epoch na to average Bob and Alice's trained models togeda and then we go use this to set the values for our global \"model\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    model.weight.set_(((alices_model.weight.data + bobs_model.weight.data) / 2).get())\n",
    "    model.bias.set_(((alices_model.bias.data + bobs_model.bias.data) / 2).get())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rinse and Repeat\n",
    "\n",
    "Now wetin we go do na to just iterate this multiple times!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterations = 10\n",
    "worker_iters = 5\n",
    "\n",
    "for a_iter in range(iterations):\n",
    "    \n",
    "    bobs_model = model.copy().send(bob)\n",
    "    alices_model = model.copy().send(alice)\n",
    "\n",
    "    bobs_opt = optim.SGD(params=bobs_model.parameters(),lr=0.1)\n",
    "    alices_opt = optim.SGD(params=alices_model.parameters(),lr=0.1)\n",
    "\n",
    "    for wi in range(worker_iters):\n",
    "\n",
    "        # We go train Bob's Model\n",
    "        bobs_opt.zero_grad()\n",
    "        bobs_pred = bobs_model(bobs_data)\n",
    "        bobs_loss = ((bobs_pred - bobs_target)**2).sum()\n",
    "        bobs_loss.backward()\n",
    "\n",
    "        bobs_opt.step()\n",
    "        bobs_loss = bobs_loss.get().data\n",
    "\n",
    "        # We go train Alice's Model\n",
    "        alices_opt.zero_grad()\n",
    "        alices_pred = alices_model(alices_data)\n",
    "        alices_loss = ((alices_pred - alices_target)**2).sum()\n",
    "        alices_loss.backward()\n",
    "\n",
    "        alices_opt.step()\n",
    "        alices_loss = alices_loss.get().data\n",
    "    \n",
    "    alices_model.move(secure_worker)\n",
    "    bobs_model.move(secure_worker)\n",
    "    with torch.no_grad():\n",
    "        model.weight.set_(((alices_model.weight.data + bobs_model.weight.data) / 2).get())\n",
    "        model.bias.set_(((alices_model.bias.data + bobs_model.bias.data) / 2).get())\n",
    "    \n",
    "    print(\"Bob:\" + str(bobs_loss) + \" Alice:\" + str(alices_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Last last, we go make sure sey our resulting model learn wella, so we go evaluate am on top test dataset.  For this proble, we go use origina data, but in practice we go use new data to understand how well the model generalizes to unseen examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = model(data)\n",
    "loss = ((preds - target) ** 2).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(preds)\n",
    "print(target)\n",
    "print(loss.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For dis toy example, the average model go dey underfitting relative to how plaintext model trained go behave, we got fit train am and we no go expose each worker's training data. We go fit add all the updated models wey we get from each worker we dey trusted aggregator togeda. Dis aggregate go prevent data leakage to the model owner.\n",
    "\n",
    "For the example wey we go do again, we go try to do our trusted aggregation directly with the gradients, if we do am like dis we go fit update the model with beta gradient estimates and we go arrive at a stronger model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Congratulations!!! - Oya Join the Community!\n",
    "\n",
    "Clap for una sef as you don finish this notebook tutorial! If you enjoy am and you wan join the movement towards privacy preserving, decentralized ownership of AI and the AI supply chain (data), follow the steps wey dey below.\n",
    "\n",
    "### Star PySyft on GitHub\n",
    "\n",
    "The easiset way to helep our community na to star the GitHub repos! This go helep raise awareness of the tools we dey build.\n",
    "\n",
    "- [Star PySyft](https://github.com/OpenMined/PySyft)\n",
    "\n",
    "### Join our Slack!\n",
    "\n",
    "To follow up bumper to bumper on how latest advancements, join our community! You can do so by filling out the form at [http://slack.openmined.org](http://slack.openmined.org)\n",
    "\n",
    "### Join a Code Project!\n",
    "\n",
    "The best way to contribute to our community na to become code contributor! You fit go to PySyft GitHub Issues page and filter for \"Projects\". E go show you all the top level Tickets giving an overview of what projects you fit join! If you no wan join any project, but you wan code small, you fit look for more \"one off\" mini-projects by searching for GitHub issues marked \"good first issue\"\n",
    "\n",
    "- [PySyft Projects](https://github.com/OpenMined/PySyft/issues?q=is%3Aopen+is%3Aissue+label%3AProject)\n",
    "- [Good First Issue Tickets](https://github.com/OpenMined/PySyft/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
    "\n",
    "### Donate\n",
    "\n",
    "If you no get time to contribute to our codebase, but still like to lend support, you fit be a Backer on our Open Collective. All donations wey we get na for our web hosting and other community expenses such as hackathons and meetups! meetups!\n",
    "\n",
    "[OpenMined's Open Collective Page](https://opencollective.com/openmined)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
